package com.lm.flight.example;

import org.apache.arrow.flight.FlightProducer;
import org.apache.arrow.flight.NoOpFlightProducer;
import org.apache.arrow.flight.Ticket;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.types.pojo.Schema;

import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.List;
/**
 * 定义一个简单的 FlightProducer，用于返回示例数据
 */
public class DemoFlightProducer extends NoOpFlightProducer {
    private final BufferAllocator allocator;

    public DemoFlightProducer(BufferAllocator allocator) {
        this.allocator = allocator;
    }

    @Override
    public void getStream(FlightProducer.CallContext context, Ticket ticket, FlightProducer.ServerStreamListener listener) {
        // 定义表结构
        Field id = new Field("id", FieldType.nullable(new ArrowType.Int(32, true)), null);
        Field name = new Field("name", FieldType.nullable(new ArrowType.Utf8()), null);
        List<Field> fields = Arrays.asList(id, name);
        Schema schema = new Schema(fields);

        // 创建 VectorSchemaRoot
        try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
            listener.start(root);

            // 获取具体类型向量并分配内存
            IntVector idVector = (IntVector) root.getVector("id");
            VarCharVector nameVector = (VarCharVector) root.getVector("name");
            root.allocateNew();

            // 写入数据（示例两行）
            idVector.setSafe(0, 1);
            nameVector.setSafe(0, "Alice".getBytes(StandardCharsets.UTF_8));

            idVector.setSafe(1, 2);
            nameVector.setSafe(1, "Bob".getBytes(StandardCharsets.UTF_8));

            // 设置有效行数（两种方式均可）
            root.setRowCount(2);
            // 或者单独设置每个向量的 value count：
            // idVector.setValueCount(2);
            // nameVector.setValueCount(2);

            // 发送数据
            listener.putNext();
            listener.completed();
        } catch (Exception e) {
            listener.error(e);
        }
    }
}
